package io.github.huayunliufeng.common.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ArrayUtils;

import javax.crypto.Cipher;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * 用于RSA加密和解密。
 *
 * @author huayu 2024/5/15 10:31
 * @version 0.0.5
 * @since 0.0.5
 */
@Slf4j
public class RsaUtils {
    /**
     * 加密算法
     */
    private static final String algorithm = "RSA";

    private Cipher cipher;

    private PublicKey publicKey;

    private PrivateKey privateKey;

    private static volatile RsaUtils rsaUtils;

    /**
     * 最大加密长度
     */
    private static final int MAX_ENCRYPT_BLOCK = 245;

    /**
     * 最大的解密长度
     */
    private static final int MAX_DECRYPT_BLOCK = 256;

    public static RsaUtils getInstance() {
        if (rsaUtils == null) {
            synchronized (RsaUtils.class) {
                if (rsaUtils == null) {
                    rsaUtils = new RsaUtils();
                }
            }
        }
        return rsaUtils;
    }

    /**
     * 相当于示例生成, 如果有其它参数可自行处理。
     *
     * @param keySize 秘钥大小, 例如2048
     */
    public static void generateRsaKeyPair(int keySize) {
        try {
            // 获取KeyPairGenerator实例并指定算法为RSA
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(algorithm);
            // 初始化密钥生成器，指定密钥长度（例如，2048位）
            keyPairGenerator.initialize(keySize);
            // 生成密钥对
            KeyPair keyPair = keyPairGenerator.generateKeyPair();
            // 获取公钥和私钥
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();
            System.out.println("Public Key: " + HylfByteUtil.base64Encoder.encodeToString(publicKey.getEncoded()));
            System.out.println("Private Key: " + HylfByteUtil.base64Encoder.encodeToString(privateKey.getEncoded()));
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
    }

    private RsaUtils() {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(algorithm);
            X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(getRsaPublicKey());
            this.publicKey = keyFactory.generatePublic(publicKeySpec);
            PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(getRsaPrivateKey());
            this.privateKey = keyFactory.generatePrivate(privateKeySpec);
            cipher = Cipher.getInstance(algorithm);
        } catch (Exception e) {
            log.error("RSAUtils初始化失败。" + e.getMessage());
        }
    }

    /**
     * 默认会在resources目录下查找rsa-private-key文件。如果有自定义的位置, 请重写该方法。
     *
     * @return 私钥的原始byte[]数组
     */
    protected byte[] getRsaPrivateKey() {
        return getRsa("rsa-private-key");
    }

    /**
     * 默认会在resources目录下查找rsa-public-key文件。如果有自定义的位置, 请重写该方法。
     *
     * @return 公钥的原始byte[]数组
     */
    protected byte[] getRsaPublicKey() {
        return getRsa("rsa-public-key");
    }

    private byte[] getRsa(String name) {
        byte[] bytes = HylfByteUtil.readAllBytes(name);
        return HylfByteUtil.base64decoder.decode(bytes);
    }

    /**
     * 解密。
     *
     * @param bytes 加密过的数据
     * @return 解密过的数据
     */
    public byte[] decodeData(byte[] bytes) {
        try {
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
        } catch (Exception ignore) {
            return bytes;
        }
        return segmentDoFinal(bytes, MAX_DECRYPT_BLOCK);
    }

    /**
     * 加密。
     *
     * @param bytes 待加密的数据
     * @return 加密过后的数据
     */
    public byte[] encodeData(byte[] bytes) {
        try {
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        } catch (Exception ignore) {
            return bytes;
        }
        return segmentDoFinal(bytes, MAX_ENCRYPT_BLOCK);
    }

    private byte[] segmentDoFinal(byte[] bytes, int limit) {
        try {
            int length = bytes.length;
            byte[] result = new byte[length * 4];
            int offset = 0, resultLength = 0;
            byte[] tmp;
            while (length > offset) {
                if (length - offset > limit) {
                    tmp = cipher.doFinal(bytes, offset, limit);
                } else {
                    tmp = cipher.doFinal(bytes, offset, length - offset);
                }
                System.arraycopy(tmp, 0, result, resultLength, tmp.length);
                offset += limit;
                resultLength += tmp.length;
            }
            return ArrayUtils.subarray(result, 0, resultLength);
        } catch (Exception e) {
            log.error("加密/解密失败。", e);
            throw new RuntimeException("加密/解密失败。");
        }
    }
}
